package protocol

import (
	"fmt"

	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/checksum"
	"gvisor.dev/gvisor/pkg/tcpip/header"
)

// 计算IPv4校验和
func CalculateIPv4Checksum(header []byte) uint16 {
	// 清除现有校验和
	header[10] = 0
	header[11] = 0

	var sum uint32

	// 按16位字计算
	for i := 0; i < len(header); i += 2 {
		if i+1 < len(header) {
			sum += uint32(header[i])<<8 + uint32(header[i+1])
		} else {
			sum += uint32(header[i]) << 8
		}
	}

	// 处理进位
	for sum>>16 > 0 {
		sum = (sum & 0xFFFF) + (sum >> 16)
	}

	return uint16(^sum)
}

func CalculatePseudoHeaderChecksumIpv6(srcIP, dstIP tcpip.Address, protocol uint8, length uint32) uint32 {
	if srcIP.Len() != 16 || dstIP.Len() != 16 {
		panic("IPv6 addresses must be 16 bytes")
	}

	var sum uint32

	srcBytes := srcIP.AsSlice()
	dstBytes := dstIP.AsSlice()

	// 源IP地址 (16字节)
	for i := 0; i < 16; i += 2 {
		sum += uint32(srcBytes[i])<<8 + uint32(srcBytes[i+1])
	}

	// 目标IP地址 (16字节)
	for i := 0; i < 16; i += 2 {
		sum += uint32(dstBytes[i])<<8 + uint32(dstBytes[i+1])
	}

	// 上层协议长度 (32位)
	sum += uint32(length >> 16)    // 高16位
	sum += uint32(length & 0xFFFF) // 低16位

	// 协议号
	sum += uint32(protocol)

	return sum
}

func CalculatePseudoHeaderChecksumIpv4(srcIP, dstIP tcpip.Address, protocol uint8, length uint16) uint32 {
	// 检查地址长度
	if srcIP.Len() != 4 || dstIP.Len() != 4 {
		panic("IPv4 addresses must be 4 bytes")
	}

	var sum uint32

	// 获取地址字节
	srcBytes := srcIP.AsSlice()
	dstBytes := dstIP.AsSlice()

	// 源IP地址 (4字节)
	sum += uint32(srcBytes[0])<<8 + uint32(srcBytes[1])
	sum += uint32(srcBytes[2])<<8 + uint32(srcBytes[3])

	// 目标IP地址 (4字节)
	sum += uint32(dstBytes[0])<<8 + uint32(dstBytes[1])
	sum += uint32(dstBytes[2])<<8 + uint32(dstBytes[3])

	// 协议号
	sum += uint32(protocol)

	// 传输层长度
	sum += uint32(length)

	return sum
}

// 计算16位校验和
func CalculateChecksum(data []byte, initialSum uint32) uint16 {
	sum := initialSum

	// 按16位字计算
	for i := 0; i < len(data); i += 2 {
		if i+1 < len(data) {
			sum += uint32(data[i])<<8 + uint32(data[i+1])
		} else {
			sum += uint32(data[i]) << 8
		}
	}

	// 处理进位
	for sum>>16 > 0 {
		sum = (sum & 0xFFFF) + (sum >> 16)
	}

	return uint16(^sum)
}

func RecalculateIcmpChecksum(pkg []byte, ipVersion uint8, srcIP, dstIP tcpip.Address) uint16 {
	// 计算新的校验和
	if ipVersion == header.IPv4Version {
		return CalculateChecksum(pkg, 0)
	}
	// ICMPv6使用伪头部校验和，协议号为58
	pseudoSum := CalculatePseudoHeaderChecksumIpv6(srcIP, dstIP, 58, uint32(len(pkg)))
	return CalculateChecksum(pkg, pseudoSum)
}

// 重新计算UDP校验和
func RecalculateTransportChecksum(pkg []byte, ipVersion uint8, src, dst tcpip.Address, protocol uint8) uint16 {

	// 计算伪头部校验和
	pseudoSum := uint32(0)
	switch ipVersion {
	case header.IPv4Version:
		pseudoSum = CalculatePseudoHeaderChecksumIpv4(src, dst, protocol, uint16(len(pkg)))
	case header.IPv6Version:
		pseudoSum = CalculatePseudoHeaderChecksumIpv6(src, dst, protocol, uint32(len(pkg)))
	default:
		panic("unsupported IP version")
	}

	// 计算UDP校验和
	checksum := CalculateChecksum(pkg, pseudoSum)

	// UDP校验和为0时表示不使用校验和（可选）
	if checksum == 0 {
		checksum = 0xFFFF
	}

	// 写入校验和
	return uint16(checksum)
}

// 更新2字节对齐地址的校验和
func ChecksumUpdate2ByteAlignedAddress(xsum uint16, old, new tcpip.Address) uint16 {
	const uint16Bytes = 2

	if old.BitLen() != new.BitLen() {
		panic(fmt.Sprintf("buffer lengths are different; old = %d, new = %d", old.BitLen()/8, new.BitLen()/8))
	}

	if oldBytes := old.BitLen() % 16; oldBytes != 0 {
		panic(fmt.Sprintf("buffer has an odd number of bytes; got = %d", oldBytes))
	}

	oldAddr := old.AsSlice()
	newAddr := new.AsSlice()

	// As per RFC 1071 page 4,
	//	(4)  Incremental Update
	//
	//        ...
	//
	//        To update the checksum, simply add the differences of the
	//        sixteen bit integers that have been changed.  To see why this
	//        works, observe that every 16-bit integer has an additive inverse
	//        and that addition is associative.  From this it follows that
	//        given the original value m, the new value m', and the old
	//        checksum C, the new checksum C' is:
	//
	//                C' = C + (-m) + m' = C + (m' - m)
	for len(oldAddr) != 0 {
		// Convert the 2 byte sequences to uint16 values then apply the increment
		// update.
		xsum = ChecksumUpdate2ByteAlignedUint16(xsum, (uint16(oldAddr[0])<<8)+uint16(oldAddr[1]), (uint16(newAddr[0])<<8)+uint16(newAddr[1]))
		oldAddr = oldAddr[uint16Bytes:]
		newAddr = newAddr[uint16Bytes:]
	}

	return xsum
}

func ChecksumUpdate2ByteAlignedUint16(xsum, old, new uint16) uint16 {
	// As per RFC 1071 page 4,
	//	(4)  Incremental Update
	//
	//        ...
	//
	//        To update the checksum, simply add the differences of the
	//        sixteen bit integers that have been changed.  To see why this
	//        works, observe that every 16-bit integer has an additive inverse
	//        and that addition is associative.  From this it follows that
	//        given the original value m, the new value m', and the old
	//        checksum C, the new checksum C' is:
	//
	//                C' = C + (-m) + m' = C + (m' - m)
	return checksum.Combine(xsum, checksum.Combine(new, ^old))
}
